import os
from PIL import Image
import numpy as np
import torch
from torch.utils.data import Dataset
from preprocess.transforms import ToTensor
from torchvision import transforms
import SimpleITK as sitk
from batchgenerators.augmentations.utils import pad_nd_image
from batchgenerators.augmentations.crop_and_pad_augmentations import crop
from preprocess.transforms import get_train_transform




# https://pytorch.org/vision/master/_modules/torchvision/transforms/functional.html
class TrainDataset(Dataset):
    def __init__(self, image_path_list,patch_size=(64,256,256),input_num=2):
        # generate image path list

        self.image_path_list = image_path_list
        self.patch_size = patch_size 
        self.input_num = input_num

    def __getitem__(self, index):
        image_path,mask_path = self.image_path_list[index]

        if self.input_num==2:
            image = sitk.GetArrayFromImage(sitk.ReadImage(image_path))
            mask = sitk.GetArrayFromImage(sitk.ReadImage(mask_path))
            input = np.stack([image,mask],0)
            input = input /255.0
        else:
            image = sitk.GetArrayFromImage(sitk.ReadImage(image_path))
            image = np.expand_dims(image,axis=0)
            input = image/255.0



        # get label from dirname
        label = int(image_path.split("/")[2])
        label = torch.tensor(label)



        

        # this will only pad patient_data if its shape is smaller than self.patch_size
        input = pad_nd_image(input, self.patch_size)
        input, patient_seg = crop(input[None],None, self.patch_size, crop_type="random")

        input = torch.tensor(input)
        if self.input_num==2:
            input = torch.squeeze(input)
        else:
            input=torch.squeeze( input , dim=0)
        
        return input.float(),label

    def __len__(self):
        return len(self.image_path_list)

    